package com.aliyun.oss.des_analysis;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Random;

public class Main {

    private static final String ALGORITHM = "DES";

    public static void main(String[] args) throws Exception {
        // Generate a key for DES
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(56); // DES key size is 56 bits
        SecretKey secretKey = keyGen.generateKey();

        // Part A: Fixed key, varying plaintext difference
        analyzeFixedKeyVaryingPlaintextDifference(secretKey);

        // Part B: Fixed plaintext, varying key difference
        analyzeFixedPlaintextVaryingKeyDifference("This is a test plaintext.".getBytes());
    }

    private static void analyzeFixedKeyVaryingPlaintextDifference(SecretKey key) throws Exception {
        byte[] plaintext1 = generateRandomPlaintext(8); // DES block size is 8 bytes
        for (int diffWeight = 1; diffWeight <= 64; diffWeight++) {
            byte[] plaintext2 = generatePlaintextWithDifferenceWeight(plaintext1, diffWeight);

            byte[] ciphertext1 = encrypt(key, plaintext1);
            byte[] ciphertext2 = encrypt(key, plaintext2);

            int hammingWeightCipher1 = hammingWeight(ciphertext1);
            int hammingWeightCipher2 = hammingWeight(ciphertext2);

            System.out.println("密钥k固定, 差分海明重量: " + diffWeight +
                    ", DES密文C1 海明重量: " + hammingWeightCipher1 +
                    ", DES密文C2 海明重量: " + hammingWeightCipher2);
        }
    }

    private static void analyzeFixedPlaintextVaryingKeyDifference(byte[] plaintext) throws Exception {
        SecretKey key1 = generateRandomKey();
        for (int diffWeight = 1; diffWeight <= 64; diffWeight++) {
            SecretKey key2 = generateKeyWithDifferenceWeight(key1, diffWeight);

            byte[] ciphertext1 = encrypt(key1, plaintext);
            byte[] ciphertext2 = encrypt(key2, plaintext);

            int hammingWeightCipher1 = hammingWeight(ciphertext1);
            int hammingWeightCipher2 = hammingWeight(ciphertext2);

            System.out.println("明文固定, 差分海明重量: " + diffWeight +
                    ", DES密文C1 海明重量: " + hammingWeightCipher1 +
                    ", DES密文C2 海明重量: " + hammingWeightCipher2);
        }
    }

    private static byte[] generateRandomPlaintext(int length) {
        byte[] plaintext = new byte[length];
        new Random().nextBytes(plaintext);
        return plaintext;
    }

    private static byte[] generatePlaintextWithDifferenceWeight(byte[] plaintext, int diffWeight) {
        byte[] plaintext2 = plaintext.clone();
        int positions = diffWeight / 8;
        int remainder = diffWeight % 8;

        Random random = new Random();
        for (int i = 0; i < positions; i++) {
            plaintext2[random.nextInt(plaintext.length)] = (byte) (random.nextInt(256));
        }
        if (remainder > 0) {
            int bytePos = random.nextInt(plaintext.length);
            int bitPos = random.nextInt(8);
            plaintext2[bytePos] = (byte) ((plaintext2[bytePos] & ~(1 << bitPos)) | (1 << bitPos));
        }
        return plaintext2;
    }

    private static SecretKey generateRandomKey() throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(56);
        return keyGen.generateKey();
    }

    private static SecretKey generateKeyWithDifferenceWeight(SecretKey key, int diffWeight) {
        byte[] keyBytes = key.getEncoded();
        byte[] newKeyBytes = keyBytes.clone();

        int positions = diffWeight / 8;
        int remainder = diffWeight % 8;

        Random random = new SecureRandom();
        for (int i = 0; i < positions; i++) {
            newKeyBytes[random.nextInt(keyBytes.length)] = (byte) (random.nextInt(256));
        }
        if (remainder > 0) {
            int bytePos = random.nextInt(keyBytes.length);
            int bitPos = random.nextInt(8);
            newKeyBytes[bytePos] = (byte) ((newKeyBytes[bytePos] & ~(1 << bitPos)) | (1 << bitPos));
        }

        return new SecretKeySpec(newKeyBytes, ALGORITHM);
    }

    private static byte[] encrypt(SecretKey key, byte[] plaintext) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        return cipher.doFinal(plaintext);
    }

    private static int hammingWeight(byte[] data) {
        int weight = 0;
        for (byte b : data) {
            weight += Integer.bitCount(b & 0xFF);
        }
        return weight;
    }
}
